import cv2
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from torch.utils.data import DataLoader
from tqdm import tqdm

from config import get_parser
from datalist import BaiDuLaneDataset
from model import DeepLab

metric_loss = 0

from torchvision import transforms

transform=transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.485, 0.456, 0.406),
                          std=(0.229, 0.224, 0.225))
])

class train():
    def __init__(self):
        self.args = get_parser()
        print(f"-----------{self.args.project_name}-------------")

        '''
        判断是否使用cuda
        '''
        use_cuda = self.args.use_cuda and torch.cuda.is_available()
        '''
        设置随机数种子，保证实验结果的可重复性
        '''

        self.device = torch.device('cuda' if use_cuda else 'cpu')
        '''
        pin_memory 锁页内存
        =True 意味着生成的Tensor数据最开始是属于内存中的锁页内存，这样将内存的Tensor转义到GPU的显存会更快一些
        主机中的内存，有两种存在方式，一是锁页，二是不锁页，
        锁页内存存放的内容在任何情况下都不会与主机的虚拟内存进行交换（注：虚拟内存就是硬盘），
        而不锁页内存在主机内存不足时，数据会存放在虚拟内存中。
        而显卡中的显存全部是锁页内存！

        当计算机的内存充足的时候，可以设置pin_memory=True。
        当系统卡住，或者交换内存使用过多的时候，设置pin_memory=False。因为pin_memory与电脑硬件性能有关，
        pytorch开发者不能确保每一个炼丹玩家都有高端设备，因此pin_memory默认为False。
        '''

        test_kwargs = {'num_workers': 0, 'pin_memory': True} if use_cuda else {}
        '''
        构造DataLoader
        '''

        self.test_dataset = BaiDuLaneDataset(root_file=self.args.data_base, phase='test')

        self.test_dataloader = DataLoader(self.test_dataset, batch_size=1, **test_kwargs)
        '''
        定义模型
        '''
        self.model = DeepLab().to(self.device)  # 这里可以替换模型 ，可用的有
        # deeplabv3+
        # deeplab_mobilenet


        '''
        根据需要加载与训练模型权重参数
        '''

        data_dict = torch.load("F:/baidu_lane_seg/weights/best.pth")['model_state_dict']
        new_data_dict = {}
        for k, v in data_dict.items():
            new_data_dict[k[7:]] = v
        self.model.load_state_dict(new_data_dict, strict=True)
        print("load pretrained model successful!")

        self.test(epoch=1)
        torch.cuda.empty_cache()
        print("model finish training")

    def test(self, epoch):
        self.model.eval()
        with torch.no_grad():

            img=Image.open(r"G:\项目\图像分割\lane_seg\ColorImage_road02\ColorImage\Record002\Camera 6\170927_063934208_Camera_6.jpg")
            img=img.crop((0,600,img.size[0],img.size[1]))
            img=img.crop((0,0,img.size[0],img.size[1]-350))
            img=img.resize((1536,512))
            temp=img
            img=transform(img).unsqueeze(0)

            outputs=self.model(img.to(self.device))
            preds=outputs.data.max(1)[1].cpu().numpy()


            pred=preds[0].astype(np.uint8)
            pred=self.test_dataset.decode_color_map(pred)

            temp=temp.convert('RGBA')
            pred=Image.fromarray(pred).convert('RGBA')

            result=Image.blend(temp,pred,0.5).convert('RGB')
            result.save('./result/'+ str(epoch) + '_result.jpg')





train = train()
